package org.september.pisces.imagesearch.service;

import java.awt.image.BufferedImage;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import org.september.core.exception.BusinessException;
import org.september.pisces.imagesearch.model.ImageEncoderModel;
import org.september.pisces.imagesearch.utils.FeatureUtils;
import org.september.pisces.imagesearch.vo.FeatureInfo;
import org.september.pisces.imagesearch.vo.ImageSearchResult;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import ai.djl.modality.cv.Image;
import ai.djl.opencv.OpenCVImageFactory;
import ai.djl.translate.TranslateException;

@Service
public class ImageFeatureService {

	private int topK = 20;
	
    private ImageEncoderModel imageEncoderModel = new ImageEncoderModel();
	
	@Autowired
	private FeatureStoreService featureStoreService;
	
	public FeatureInfo addImage(String filePath) {
		checkInit();
		FeatureInfo info = new FeatureInfo();
		try {
			Path imageFile = Paths.get(filePath);
			info.setData(extractFeatures(imageFile));
			info.setPath(filePath);
			featureStoreService.add(info);
		}catch(Exception ex) {
			ex.printStackTrace();
		}
		return info;
	}
	
	private float[] extractFeatures(Path filePath) throws IOException, TranslateException {
		Image image = OpenCVImageFactory.getInstance().fromFile(filePath);
        float[] feature = imageEncoderModel.predict(image);
        return feature;
	}
	
	private float[] extractFeatures(BufferedImage bi) throws IOException, TranslateException {
		Image image = OpenCVImageFactory.getInstance().fromImage(bi);
        float[] feature = imageEncoderModel.predict(image);
        return feature;
	}
	
	public List<ImageSearchResult> search(String filePath) throws TranslateException, IOException {
		checkInit();
		Path imageFile = Paths.get(filePath);
		Image image = OpenCVImageFactory.getInstance().fromFile(imageFile);
        float[] vectorToSearch = imageEncoderModel.predict(image);
        
        List<ImageSearchResult> list = new ArrayList<>();
        for(FeatureInfo info : featureStoreService.getFeatures()) {
        	float score = FeatureUtils.cosineSim(vectorToSearch, info.getData());
        	ImageSearchResult res = new ImageSearchResult();
        	res.setScore(score);
        	res.setFileId(info.getPath());
        	list.add(res);
        }
        Collections.sort(list);
        if (topK >= list.size()) {
            return list;
        } else {
        	list = list.subList(0, topK);
            return list;
        }
	}
	
	private void checkInit() {
		if(imageEncoderModel.getInitFlag()==false) {
			try {
				String path = this.getClass().getResource("./CLIP-ViT-B-32-IMAGE.pt").getPath();
				path = path.substring(1);
				imageEncoderModel.init(path, 4);
			} catch (Exception e) {
				throw new BusinessException("加载图片搜索模型失败",e);
			}
		}
	}
}
